import os
import json
import cv2

def read_box_trace(boxes_path):
    frame_2_boxes = {}
    with open(boxes_path, 'r') as fin:
        for line in fin.readlines():
            frame_idx,box_str = line.strip().split(':')
            frame_idx = int(frame_idx)+1
            boxes = json.loads(box_str)
            frame_2_boxes[frame_idx] = boxes
    return frame_2_boxes

def box_overlap_ratio(box1, box2, which_box=1):
    """
    计算重叠面积占指定矩形框的比例
    
    参数:
    box1 -- 第一个矩形框
    box2 -- 第二个矩形框
    which_box -- 1表示计算占box1的比例，2表示占box2的比例
    
    返回:
    ratio -- 重叠面积占指定矩形框的比例
    """
    # 计算相交区域
    x_left = max(box1[0], box2[0])
    y_top = max(box1[1], box2[1])
    x_right = min(box1[2], box2[2])
    y_bottom = min(box1[3], box2[3])
    
    if x_right < x_left or y_bottom < y_top:
        return 0.0
    
    intersection_area = (x_right - x_left) * (y_bottom - y_top)
    
    if which_box == 1:
        box_area = (box1[2] - box1[0]) * (box1[3] - box1[1])
    else:
        box_area = (box2[2] - box2[0]) * (box2[3] - box2[1])
    
    return intersection_area / box_area

def box_in_large_top_area(small_box, large_box):
    """
    计算box1在box2的上半部分
    参数:
        small_box : 小框坐标
        large_box : 大框坐标
    返回:
        bool : True 如果小框位于大框的上半部分,否则False
    """
    lx1, ly1, lx2, ly2 = large_box
    sx1, sy1, sx2, sy2 = small_box

    large_mid_y = (ly1 + ly2)/2
    
    # 检查小框是否完全在大框内
    if not (lx1 <= sx1 and lx2 >= sx2 and ly1 <= sy1 and ly2 >= sy2):
        return False
    
    # 检查小框是否完全在上半部分（小框的底部不超过大框的中点）
    return sy2 <= large_mid_y

def cnt_person_num(person_2_number_list):
    person_2_num = {}
    for person_id, number_list in person_2_number_list.items():
        num_2_cnt = {}
        for num in number_list:
            cnt = num_2_cnt.get(num, 0)
            cnt += 1
            num_2_cnt[num] = cnt
        sorted_num_cnt = sorted(num_2_cnt.items(), key=lambda kv:kv[1], reverse=True)
        num_max = sorted_num_cnt[0][0]
        person_2_num[person_id] = num_max 
    return person_2_num

def apply_num_to_trace(frame_2_person_boxes, person_2_number):
    frame_2_out = {}
    for frame_id in frame_2_person_boxes:
        person_used = {}
        person_boxes = frame_2_person_boxes[frame_id]
        out_person_boxes = []
        for person_box in person_boxes:
            is_person_used = False
            person_id, p_x1, p_y1, p_x2, p_y2 = person_box
            for p in person_used:
                p_1 = (p_x1, p_y1, p_x2, p_y2)
                if box_overlap_ratio(p_1, p) > 0.9:
                    is_person_used = True
                    break
            if not is_person_used:
                person_used[(p_x1, p_y1, p_x2, p_y2)] = True
                num = person_2_number.get(person_id, None)
                out_person_boxes.append((person_id, num, p_x1, p_y1, p_x2, p_y2))
        frame_2_out[frame_id] = out_person_boxes
    return frame_2_out

def write_trace_out(out_path, frame_2_out):
    with open(out_path, 'w') as fout:
        for frame_idx in frame_2_out:
            person_with_num = frame_2_out.get(frame_idx, [])
            out_persons = json.dumps(person_with_num)
            line = "{}:{}\n".format(frame_idx, out_persons)
            fout.write(line)

def draw_box(frame, out_boxes):
    font = cv2.FONT_HERSHEY_COMPLEX
    color = (255, 0, 0)
    for o_box in out_boxes:
        #print("===== draw box ", o_box)
        person_id, num, p_x1, p_y1, p_x2, p_y2 = o_box
        cv2.rectangle(frame, (p_x1, p_y1), (p_x2, p_y2), color, 5)
        cv2.putText(frame, str(num), (p_x1-10, p_y1-10), font, 3, color, 5) 
    return frame

def draw_trace_out(in_video_path, out_video_path, frame_2_out):
    more = True
    frame_idx = -1
    vid_reader = cv2.VideoCapture(in_video_path)
    width = int(vid_reader.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(vid_reader.get(cv2.CAP_PROP_FRAME_HEIGHT)) 
    vid_writer = cv2.VideoWriter(out_video_path, cv2.VideoWriter_fourcc(*"mp4v"), 25, (width, height))
    
    while more:
        more, frame = vid_reader.read()
        frame_idx += 1
        if frame_idx % 10 == 0:
            print("== finish draw frame ", frame_idx) 
        if not more:
            break
        out_boxes = frame_2_out.get(frame_idx, [])
        frame = draw_box(frame, out_boxes)
        vid_writer.write(frame)
    vid_reader.release()
    vid_writer.release()

def extract_person_2_num(frame_2_person_boxes, frame_2_num_boxes):
    person_2_number_list = {}
    for frame_idx, person_boxes in frame_2_person_boxes.items():
        num_boxes = frame_2_num_boxes.get(frame_idx, [])
        for person_box in person_boxes:
            person_id, p_x1, p_y1, p_x2, p_y2 = person_box
            p_box = [p_x1, p_y1, p_x2, p_y2] 
            for num_box in num_boxes:
                number_box, num = num_box
                number_cross_ratio = box_overlap_ratio(number_box, p_box)
                in_top = box_in_large_top_area(number_box, p_box)

                if number_cross_ratio > 0.95 and in_top:
                    number_list = person_2_number_list.get(person_id, [])
                    number_list.append(num)
                    person_2_number_list[person_id] = number_list
                #print("frame : {} person_Id {} box ratio {} for player number {}".format(frame_idx, person_id, number_cross_ratio, num))
    #person_2_num = cnt_person_num(person_2_number_list)
    #return person_2_num
    return person_2_number_list

def merge_person_2_number(person_2_num1, person_2_num2):
    person_2_num_out = {}
    for k,v in person_2_num1.items():
        if k in person_2_num2:
            v.extend(person_2_num2[k])
        person_2_num_out[k] = v
    for k,v in person_2_num2.items():
        if k not in person_2_num_out:
            person_2_num_out[k] = v
    return person_2_num_out 

def track_trace():
    frame_2_bboxes_path = '/home/imvision/hy/samurai-master/data/bboxes_000.txt'
    frame_2_number_path = './datas/v000_number_out.json'

    frame_2_person_boxes = read_box_trace(frame_2_bboxes_path)
    frame_2_num_boxes = read_box_trace(frame_2_number_path)
    person_2_num1 = extract_person_2_num(frame_2_person_boxes, frame_2_num_boxes)
    print("======= person 2 number ",person_2_num1)

    frame_2_boxes_path2 =  '/home/imvision/hy/samurai-master/data/bboxes_001.txt'
    frame_2_number_path2 = './datas/v001_number_out.json'
    frame_2_person_boxes2 = read_box_trace(frame_2_boxes_path2)
    frame_2_num_boxes2 = read_box_trace(frame_2_number_path2)
    person_2_num2 = extract_person_2_num(frame_2_person_boxes2, frame_2_num_boxes2)
    print("====== person 2 number 2 ", person_2_num2)
    #person_2_num = {0:'6'}
    person_2_num_out = merge_person_2_number(person_2_num1, person_2_num2)
    person_2_num = cnt_person_num(person_2_num_out)

    vid = 9
    frame_2_cur_boxes_path = '/home/imvision/hy/samurai-master/data/bboxes_%03d.txt' % vid
    in_video_path = '/home/imvision/hy/samurai-master/data/output%03d.mp4' % vid
    out_video_path = './datas/v%03d_person_all_number_view.mp4' % vid
    frame_2_cur_boxes = read_box_trace(frame_2_cur_boxes_path)

    print("====== cnt person for number ", person_2_num)
    frame_2_out = apply_num_to_trace(frame_2_cur_boxes, person_2_num)
    draw_trace_out(in_video_path, out_video_path, frame_2_out)

def read_person_trace(json_dir):
    frame_2_person = {}
    for idx, json_file in enumerate(os.listdir(json_dir)):
        frame_idx = int(json_file.strip('.json').split('_')[1])
        json_path = os.path.join(json_dir, json_file)
        with open(json_path, 'r') as fin:
            frame_boxes = json.loads(fin.read())['labels']
        frame_2_person[frame_idx] = frame_boxes
    return frame_2_person

def compute_player_2_num(grouned_sam2_json_dir, number_json, thresh=0.95):
        
    frame_2_number_boxes = read_box_trace(number_json)    

    frame_2_person_boxes = read_person_trace(grouned_sam2_json_dir)

    player_id_2_numbers = {}
    for frame_idx in range(1, 5000):
        number_boxes = frame_2_number_boxes.get(frame_idx, [])
        for number_box_number in number_boxes:
            number_box = number_box_number[0]
            number = number_box_number[1]
            score = number_box_number[2]
            if score >= thresh:
                person_id_2_boxes = frame_2_person_boxes.get(frame_idx, {})
                print("=== number boxes ", number_box, " with number ", number, " score ", score)
                #print("=== person boxes ", person_id_2_boxes)
                for player in person_id_2_boxes:
                    player_info = person_id_2_boxes[player]
                    player_id = player_info['instance_id']
                    x1, y1, x2, y2 = player_info['x1'], player_info['y1'], player_info['x2'], player_info['y2']
                    person_box = (x1, y1, x2, y2)
                    overlap_ratio = box_overlap_ratio(number_box, person_box)
                    in_top = box_in_large_top_area(number_box, person_box)
                    print("=== player_id ", player_id, " number ", number, " overlap ", overlap_ratio, " in_top ", in_top)
                    if overlap_ratio > 0.95 and in_top:
                        print("=== player ", player_id, "overlap ", overlap_ratio)
                        number_list = player_id_2_numbers.get(player_id, []) 
                        number_list.append(number)
                        player_id_2_numbers[player_id] = number_list
    #print(player_id_2_numbers)             
    player_2_num = cnt_person_num(player_id_2_numbers)
    print(player_2_num)
    return player_2_num

if __name__ == '__main__':
    grouned_sam2_json_dir = '/home/imvision/hy/Grounded-SAM-2/outputs_cba/json_data'
    number_json = './datas/v0_to_20_number_out2.json'

    player_2_num = compute_player_2_num(grouned_sam2_json_dir, number_json)
    print(player_2_num)    

    in_video_path = '/home/imvision/hy/samurai-master/data/output_0_to_20.mp4'
    out_video_path = './datas/v0_to_20_number_trace_no_repeat2.mp4'
    frame_2_person_boxes = read_person_trace(grouned_sam2_json_dir)

    more = True
    frame_idx = 0
    vid_reader = cv2.VideoCapture(in_video_path)
    width = int(vid_reader.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(vid_reader.get(cv2.CAP_PROP_FRAME_HEIGHT)) 
    vid_writer = cv2.VideoWriter(out_video_path, cv2.VideoWriter_fourcc(*"mp4v"), 25, (width, height))
    
    while more:
        more, frame = vid_reader.read()
        frame_idx += 1
        frame_number_exist = {}
        if frame_idx % 10 == 0:
            print("== finish draw frame ", frame_idx) 
        if not more:
            break
        if frame_idx >= 5000:
            break
        person_boxes = frame_2_person_boxes.get(frame_idx, [])
        out_boxes = []
        for p_id, person_box in person_boxes.items():
            instance_id = person_box['instance_id']
            x1 = person_box['x1']
            y1 = person_box['y1']
            x2 = person_box['x2']
            y2 = person_box['y2']
            if instance_id not in player_2_num:
                print("== invalid instance ", instance_id, " in frame ", frame_idx)
                continue
            player_num_out = player_2_num[instance_id]
            if player_num_out in frame_number_exist:
                print("====== number been used ", player_num_out)
                continue
            frame_number_exist[player_num_out] = True
            out_box = [instance_id, player_num_out, x1, y1, x2, y2]
            out_boxes.append(out_box)
        frame = draw_box(frame, out_boxes)

        font = cv2.FONT_HERSHEY_COMPLEX
        color = (255, 0, 0)
        cv2.putText(frame, str(frame_idx), (10, 100), font, 3, color, 5) 
        vid_writer.write(frame)
    vid_reader.release()
    vid_writer.release()
